// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <type_traits>

#include <melon/container/f14_map.h>
#include <melon/hash/hash.h>
#include <turbo/log/logging.h>

#include <pollux/type/type.h>
#include <pollux/vector/base_vector.h>
#include <pollux/vector/lazy_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    class RowVector : public BaseVector {
    public:
        RowVector(const RowVector &) = delete;

        RowVector &operator=(const RowVector &) = delete;

        RowVector(
            pollux::memory::MemoryPool *pool,
            std::shared_ptr<const Type> type,
            BufferPtr nulls,
            size_t length,
            std::vector<VectorPtr> children,
            std::optional<vector_size_t> nullCount = std::nullopt)
            : BaseVector(
                  pool,
                  type,
                  VectorEncoding::Simple::ROW,
                  nulls,
                  length,
                  std::nullopt,
                  nullCount,
                  1),
              childrenSize_(children.size()),
              children_(std::move(children)) {
            // Some columns may not be projected out
            POLLUX_CHECK_LE(children_.size(), type->size());
            [[maybe_unused]] const auto *rowType =
                    dynamic_cast<const RowType *>(type.get());

            // Check child vector types.
            // This can be an expensive operation, so it's only done at debug time.
            for (auto i = 0; i < children_.size(); i++) {
                const auto &child = children_[i];
                if (child) {
                    POLLUX_DCHECK(
                        child->type()->kindEquals(type->childAt(i)),
                        "Got type {} for field `{}` at position {}, but expected {}.",
                        child->type()->toString(),
                        rowType->nameOf(i),
                        i,
                        type->childAt(i)->toString());
                    BaseVector::in_memory_bytes_ += child->in_memory_bytes();
                }
            }
            updateContainsLazyNotLoaded();
        }

        static std::shared_ptr<RowVector> createEmpty(
            std::shared_ptr<const Type> type,
            pollux::memory::MemoryPool *pool);

        virtual ~RowVector() override {
        }

        bool contains_null_at(vector_size_t idx) const override;

        std::optional<int32_t> compare(
            const BaseVector *other,
            vector_size_t index,
            vector_size_t otherIndex,
            CompareFlags flags) const override;

        uint64_t hash_value_at(vector_size_t index) const override;

        BaseVector *loaded_vector() override;

        const BaseVector *loaded_vector() const override {
            return const_cast<RowVector *>(this)->loaded_vector();
        }

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        /// Return the number of child vectors.
        /// This will exactly match the number of fields.
        size_t childrenSize() const {
            return childrenSize_;
        }

        /// Resize a row vector by adding trailing nulls to the top level row without
        /// resizing children.
        /// Caller should ensure that the vector is unique before calling this method.
        void appendNulls(vector_size_t numberOfRows);

        /// Get the child vector at a given offset.
        VectorPtr &childAt(column_index_t index) {
            POLLUX_CHECK_LT(
                index,
                childrenSize_,
                "Trying to access non-existing child in RowVector: {}",
                toString());
            return children_[index];
        }

        const VectorPtr &childAt(column_index_t index) const {
            POLLUX_CHECK_LT(
                index,
                childrenSize_,
                "Trying to access non-existing child in RowVector: {}",
                toString());
            return children_[index];
        }

        /// Returns child vector for the specified field name. Throws if field with
        /// specified name doesn't exist.
        VectorPtr &childAt(const std::string &name) {
            return children_[type_->as_row().getChildIdx(name)];
        }

        const VectorPtr &childAt(const std::string &name) const {
            return children_[type_->as_row().getChildIdx(name)];
        }

        std::vector<VectorPtr> &children() {
            return children_;
        }

        const std::vector<VectorPtr> &children() const {
            return children_;
        }

        void set_type(const TypePtr &type) override;

        void copy(
            const BaseVector *source,
            vector_size_t targetIndex,
            vector_size_t sourceIndex,
            vector_size_t count) override;

        void copy(
            const BaseVector *source,
            const SelectivityVector &rows,
            const vector_size_t *toSourceRow) override;

        void copy_ranges(
            const BaseVector *source,
            const melon::Range<const CopyRange *> &ranges) override;

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            std::vector<VectorPtr> copiedChildren(children_.size());

            for (auto i = 0; i < children_.size(); ++i) {
                copiedChildren[i] = children_[i]->copy_preserve_encodings(pool);
            }

            auto selfPool = pool ? pool : pool_;
            return std::make_shared<RowVector>(
                selfPool,
                type_,
                AlignedBuffer::copy(selfPool, nulls_),
                length_,
                copiedChildren,
                nullCount_);
        }

        uint64_t retained_size() const override {
            auto size = BaseVector::retained_size();
            for (auto &child: children_) {
                if (child) {
                    size += child->retained_size();
                }
            }
            return size;
        }

        uint64_t estimate_flat_size() const override;

        using BaseVector::toString;

        std::string toString(vector_size_t index) const override;

        void ensure_writable(const SelectivityVector &rows) override;

        bool is_writable() const override;

        /// Calls BaseVector::prepare_for_reuse() to check and reset nulls buffer if
        /// needed, then calls BaseVector::prepare_for_reuse(child, 0) for all children.
        void prepare_for_reuse() override;

        bool may_have_nulls_recursive() const override {
            if (BaseVector::may_have_nulls_recursive()) {
                return true;
            }

            for (const auto &child: children_) {
                if (child->may_have_nulls_recursive()) {
                    return true;
                }
            }

            return false;
        }

        VectorPtr slice(vector_size_t offset, vector_size_t length) const override;

        bool containsLazyNotLoaded() const {
            return containsLazyNotLoaded_;
        }

        void updateContainsLazyNotLoaded() const;

        void validate(const VectorValidateOptions &options) const override;

        /// Only calls BaseVector::resize and doesnt resize the children.
        /// This function is present for backwards compatibility,
        /// until the few places that require it are migrated over.
        void unsafeResize(vector_size_t newSize, bool setNotNull = true);

        /// Resizes the parent row container and also recursively resizes the
        /// children. Note that this function will throw if the children are not
        /// uniquely referenced by the parent when increasing the size.
        /// Note : If the child is null, then it will stay null after the resize.
        void resize(vector_size_t newSize, bool setNotNull = true) override;

        /// Push all dictionary encoding to the leaf vectors of a RowVector tree
        /// (i.e. we traverse the tree consists of RowVectors, possibly wrapped in
        /// DictionaryVector, and no traverse into other complex types like array or
        /// map children).  When the input is not ROW, we combine adjacent DICT
        /// layers.
        ///
        /// When new nulls are introduced on RowVector, we combine it
        /// with the existing nulls on RowVector.  The input vector should not contain
        /// any unloaded lazy.
        ///
        /// This is used for example in writing Nimble ArrayWithOffsets and
        /// SlidingWindowMap.
        static VectorPtr pushDictionaryToRowVectorLeaves(const VectorPtr &input);

        VectorPtr &rawVectorForBatchReader() {
            return rawVectorForBatchReader_;
        }

    private:
        vector_size_t childSize() const {
            bool allConstant = false;
            for (auto &child: children_) {
                if (child) {
                    if (!child->is_constant_encoding()) {
                        return child->size();
                    }
                    allConstant = true;
                }
            }
            if (!allConstant) {
                // All are nullptr.
                return 0;
            }
            // If all children are constants which do not have a meaningful
            // size, the size is one past the last referenced child index.
            return BaseVector::length_;
        }

        void appendToChildren(
            const RowVector *source,
            vector_size_t sourceIndex,
            vector_size_t count,
            vector_size_t childSize);

        const size_t childrenSize_;
        mutable std::vector<VectorPtr> children_;

        // Flag to indicate if any children of this vector contain lazy vector that
        // has not been loaded.  Used to optimize recursive laziness check.  This will
        // be initialized in the constructor, and should be updated by calling
        // updateContainsLazyNotLoaded whenever a new lazy child is set (e.g. in table
        // scan), or a lazy child is loaded (e.g. in LazyVector::ensureLoadedRows and
        // loaded_vector).
        mutable bool containsLazyNotLoaded_;

        // Flag to indicate all children has been loaded (non-recursively).  Used to
        // optimize loaded_vector calls.  If this is true, we don't recurse into
        // children to check if they are loaded.  Will be set to true when
        // loaded_vector is called, and reset to false when updateContainsLazyNotLoaded
        // is called (i.e. some children are likely updated to lazy).
        mutable bool childrenLoaded_ = false;

        // For some non-selective reader, we need to keep the original vector that is
        // unprojected and unfilterd, and reuse its memory.
        VectorPtr rawVectorForBatchReader_;
    };

    /// Common parent class for ARRAY and MAP vectors.  Contains 'offsets' and
    /// 'sizes' data and provide manipulations on them.
    struct ArrayVectorBase : BaseVector {
        ArrayVectorBase(const ArrayVectorBase &) = delete;

        const BufferPtr &offsets() const {
            return offsets_;
        }

        const BufferPtr &sizes() const {
            return sizes_;
        }

        const vector_size_t *rawOffsets() const {
            return rawOffsets_;
        }

        const vector_size_t *rawSizes() const {
            return rawSizes_;
        }

        vector_size_t offsetAt(vector_size_t index) const {
            return rawOffsets_[index];
        }

        vector_size_t sizeAt(vector_size_t index) const {
            return rawSizes_[index];
        }

        BufferPtr mutableOffsets(size_t size) {
            BaseVector::resize_indices(length_, size, pool_, offsets_, &rawOffsets_);
            return offsets_;
        }

        BufferPtr mutableSizes(size_t size) {
            BaseVector::resize_indices(length_, size, pool_, sizes_, &rawSizes_);
            return sizes_;
        }

        void resize(vector_size_t size, bool setNotNull = true) override {
            if (BaseVector::length_ < size) {
                BaseVector::resize_indices(length_, size, pool_, offsets_, &rawOffsets_);
                BaseVector::resize_indices(length_, size, pool_, sizes_, &rawSizes_);
            }
            BaseVector::resize(size, setNotNull);
        }

        /// Its the caller responsibility to make sure that `offsets_` and `sizes_`
        /// are safe to write at index i, i.ex not shared, or not large enough.
        void
        setOffsetAndSize(vector_size_t i, vector_size_t offset, vector_size_t size) {
            DKCHECK_LT(i, BaseVector::length_);
            offsets_->asMutable<vector_size_t>()[i] = offset;
            sizes_->asMutable<vector_size_t>()[i] = size;
        }

        /// Check if there is any overlapping [offset, size] ranges.
        bool hasOverlappingRanges() const {
            std::vector<vector_size_t> indices;
            return hasOverlappingRanges(
                size(), raw_nulls(), rawOffsets_, rawSizes_, indices);
        }

        /// Check if there is any overlapping [offset, size] ranges for any non-null
        /// non-empty rows.
        ///
        /// When return true (overlap exists), `indices' will be populated with
        /// ARRAY/MAP indices sorted by offsets.
        static bool hasOverlappingRanges(
            vector_size_t size,
            const uint64_t *nulls,
            const vector_size_t *offsets,
            const vector_size_t *sizes,
            std::vector<vector_size_t> &indices);

    protected:
        ArrayVectorBase(
            pollux::memory::MemoryPool *pool,
            std::shared_ptr<const Type> type,
            VectorEncoding::Simple encoding,
            BufferPtr nulls,
            size_t length,
            std::optional<vector_size_t> nullCount,
            BufferPtr offsets,
            BufferPtr lengths)
            : BaseVector(
                  pool,
                  type,
                  encoding,
                  std::move(nulls),
                  length,
                  std::nullopt /*distinctValueCount*/,
                  nullCount),
              offsets_(std::move(offsets)),
              rawOffsets_(offsets_->as<vector_size_t>()),
              sizes_(std::move(lengths)),
              rawSizes_(sizes_->as<vector_size_t>()) {
            POLLUX_CHECK_GE(
                offsets_->capacity(), BaseVector::length_ * sizeof(vector_size_t));
            POLLUX_CHECK_GE(
                sizes_->capacity(), BaseVector::length_ * sizeof(vector_size_t));
        }

        void copyRangesImpl(
            const BaseVector *source,
            const melon::Range<const CopyRange *> &ranges,
            VectorPtr *targetValues,
            VectorPtr *targetKeys);

        void validateArrayVectorBase(
            const VectorValidateOptions &options,
            vector_size_t minChildVectorSize) const;

    protected:
        BufferPtr offsets_;
        const vector_size_t *rawOffsets_;
        BufferPtr sizes_;
        const vector_size_t *rawSizes_;
    };

    class ArrayVector : public ArrayVectorBase {
    public:
        ArrayVector(const ArrayVector &) = delete;

        ArrayVector &operator=(const ArrayVector &) = delete;

        ArrayVector(
            pollux::memory::MemoryPool *pool,
            std::shared_ptr<const Type> type,
            BufferPtr nulls,
            size_t length,
            BufferPtr offsets,
            BufferPtr lengths,
            VectorPtr elements,
            std::optional<vector_size_t> nullCount = std::nullopt)
            : ArrayVectorBase(
                  pool,
                  type,
                  VectorEncoding::Simple::ARRAY,
                  std::move(nulls),
                  length,
                  nullCount,
                  std::move(offsets),
                  std::move(lengths)),
              elements_(BaseVector::get_or_create_empty(
                  std::move(elements),
                  type->childAt(0),
                  pool)) {
            POLLUX_CHECK_EQ(type->kind(), TypeKind::ARRAY);
            POLLUX_CHECK(
                elements_->type()->kindEquals(type->childAt(0)),
                "Unexpected element type: {}. Expected: {}",
                elements_->type()->toString(),
                type->childAt(0)->toString());
        }

        bool contains_null_at(vector_size_t idx) const override;

        std::optional<int32_t> compare(
            const BaseVector *other,
            vector_size_t index,
            vector_size_t otherIndex,
            CompareFlags flags) const override;

        uint64_t hash_value_at(vector_size_t index) const override;

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        const VectorPtr &elements() const {
            return elements_;
        }

        VectorPtr &elements() {
            return elements_;
        }

        void setElements(VectorPtr elements) {
            elements_ = BaseVector::get_or_create_empty(
                std::move(elements), type()->childAt(0), pool_);
        }

        void set_type(const TypePtr &type) override;

        void copy_ranges(
            const BaseVector *source,
            const melon::Range<const CopyRange *> &ranges) override;

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : pool_;
            return std::make_shared<ArrayVector>(
                selfPool,
                type_,
                AlignedBuffer::copy(selfPool, nulls_),
                length_,
                AlignedBuffer::copy(selfPool, offsets_),
                AlignedBuffer::copy(selfPool, sizes_),
                elements_->copy_preserve_encodings(pool),
                nullCount_);
        }

        uint64_t retained_size() const override {
            return BaseVector::retained_size() + offsets_->capacity() +
                   sizes_->capacity() + elements_->retained_size();
        }

        uint64_t estimate_flat_size() const override;

        using BaseVector::toString;

        std::string toString(vector_size_t index) const override;

        void ensure_writable(const SelectivityVector &rows) override;

        bool is_writable() const override;

        /// Calls BaseVector::prepare_for_reuse() to check and reset nulls buffer if
        /// needed, checks and resets offsets and sizes buffers, zeros out offsets and
        /// sizes if reusable, calls BaseVector::prepare_for_reuse(elements, 0) for the
        /// elements vector.
        void prepare_for_reuse() override;

        bool may_have_nulls_recursive() const override {
            return BaseVector::may_have_nulls_recursive() ||
                   elements_->may_have_nulls_recursive();
        }

        VectorPtr slice(vector_size_t offset, vector_size_t length) const override;

        void validate(const VectorValidateOptions &options) const override;

    private:
        VectorPtr elements_;
    };

    class MapVector : public ArrayVectorBase {
    public:
        MapVector(const MapVector &) = delete;

        MapVector &operator=(const MapVector &) = delete;

        MapVector(
            pollux::memory::MemoryPool *pool,
            std::shared_ptr<const Type> type,
            BufferPtr nulls,
            size_t length,
            BufferPtr offsets,
            BufferPtr sizes,
            VectorPtr keys,
            VectorPtr values,
            std::optional<vector_size_t> nullCount = std::nullopt,
            bool sortedKeys = false)
            : ArrayVectorBase(
                  pool,
                  type,
                  VectorEncoding::Simple::MAP,
                  std::move(nulls),
                  length,
                  nullCount,
                  std::move(offsets),
                  std::move(sizes)),
              keys_(BaseVector::get_or_create_empty(
                  std::move(keys),
                  type->childAt(0),
                  pool)),
              values_(BaseVector::get_or_create_empty(
                  std::move(values),
                  type->childAt(1),
                  pool)),
              sortedKeys_{sortedKeys} {
            POLLUX_CHECK_EQ(type->kind(), TypeKind::MAP);

            POLLUX_CHECK(
                keys_->type()->kindEquals(type->childAt(0)),
                "Unexpected key type: {}. Expected: {}",
                keys_->type()->toString(),
                type->childAt(0)->toString());
            POLLUX_CHECK(
                values_->type()->kindEquals(type->childAt(1)),
                "Unexpected value type: {}. Expected: {}",
                values_->type()->toString(),
                type->childAt(1)->toString());
        }

        virtual ~MapVector() override {
        }

        bool contains_null_at(vector_size_t idx) const override;

        std::optional<int32_t> compare(
            const BaseVector *other,
            vector_size_t index,
            vector_size_t otherIndex,
            CompareFlags flags) const override;

        uint64_t hash_value_at(vector_size_t index) const override;

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        const VectorPtr &mapKeys() const {
            return keys_;
        }

        VectorPtr &mapKeys() {
            return keys_;
        }

        const VectorPtr &mapValues() const {
            return values_;
        }

        VectorPtr &mapValues() {
            return values_;
        }

        void set_type(const TypePtr &type) override;

        bool hasSortedKeys() const {
            return sortedKeys_;
        }

        void setKeysAndValues(VectorPtr keys, VectorPtr values) {
            keys_ = BaseVector::get_or_create_empty(
                std::move(keys), type()->childAt(0), pool_);
            values_ = BaseVector::get_or_create_empty(
                std::move(values), type()->childAt(1), pool_);
        }

        void copy_ranges(
            const BaseVector *source,
            const melon::Range<const CopyRange *> &ranges) override;

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : pool_;
            return std::make_shared<MapVector>(
                selfPool,
                type_,
                AlignedBuffer::copy(selfPool, nulls_),
                length_,
                AlignedBuffer::copy(selfPool, offsets_),
                AlignedBuffer::copy(selfPool, sizes_),
                keys_->copy_preserve_encodings(pool),
                values_->copy_preserve_encodings(pool),
                nullCount_,
                sortedKeys_);
        }

        uint64_t retained_size() const override {
            return BaseVector::retained_size() + offsets_->capacity() +
                   sizes_->capacity() + keys_->retained_size() + values_->retained_size();
        }

        uint64_t estimate_flat_size() const override;

        using BaseVector::toString;

        std::string toString(vector_size_t index) const override;

        /// Sorts all maps smallest key first. This enables linear time
        /// comparison and log time lookup.  This may only be done if there
        /// are no other references to 'map'. Checks that 'map' is uniquely
        /// referenced. This is guaranteed after construction or when
        /// retrieving values from aggregation or join row containers.
        static void canonicalize(
            const std::shared_ptr<MapVector> &map,
            bool useStableSort = false);

        /// Returns indices into the map at 'index' such
        /// that keys[indices[i]] < keys[indices[i + 1]].
        std::vector<vector_size_t> sortedKeyIndices(vector_size_t index) const;

        void ensure_writable(const SelectivityVector &rows) override;

        bool is_writable() const override;

        /// Calls BaseVector::prepare_for_reuse() to check and reset nulls buffer if
        /// needed. Checks and re-allocate offsets and sizes buffers to have
        /// BaseVector::length_, or zeros out offsets and sizes if reusable, calls
        /// BaseVector::prepare_for_reuse(keys|values, 0) for the keys and values
        /// vectors.
        void prepare_for_reuse() override;

        bool may_have_nulls_recursive() const override {
            return BaseVector::may_have_nulls_recursive() ||
                   keys_->may_have_nulls_recursive() || values_->may_have_nulls_recursive();
        }

        VectorPtr slice(vector_size_t offset, vector_size_t length) const override;

        void validate(const VectorValidateOptions &options) const override;

        /// Update this map vector (base) with a list of map vectors (updates) of same
        /// size.  Maps are updated row-wise, i.e. for a certain key in each row, we
        /// keep the entry from the last update map containing the key.  If no update
        /// map contains the key, we use the entry from base.  Any null map in either
        /// base or updates creates a null row in the result.
        std::shared_ptr<MapVector> update(
            const std::vector<std::shared_ptr<MapVector> > &others) const;

        /// Same as the other `update' but can handle encodings on inputs.
        std::shared_ptr<MapVector> update(
            const melon::Range<DecodedVector *> &others) const;

    protected:
        virtual void reset_data_dependent_flags(const SelectivityVector *rows) override {
            BaseVector::reset_data_dependent_flags(rows);
            sortedKeys_ = false;
        }

    private:
        // Returns true if the keys for map at 'index' are sorted from first
        // to last in the type's collation order.
        bool isSorted(vector_size_t index) const;

        // makes a Buffer with 0, 1, 2,... size-1. This is later sorted to
        // get elements in key order in each map.
        BufferPtr elementIndices() const;

        template<TypeKind kKeyTypeKind>
        std::shared_ptr<MapVector> updateImpl(
            const melon::Range<DecodedVector *> &others) const;

        VectorPtr keys_;
        VectorPtr values_;
        bool sortedKeys_;
    };

    using RowVectorPtr = std::shared_ptr<RowVector>;
    using ArrayVectorPtr = std::shared_ptr<ArrayVector>;
    using MapVectorPtr = std::shared_ptr<MapVector>;

    // Allocates a buffer to fit at least 'size' offsets and initializes them to
    // zero.
    inline BufferPtr allocateOffsets(vector_size_t size, memory::MemoryPool *pool) {
        return AlignedBuffer::allocate<vector_size_t>(size, pool, 0);
    }

    // Allocates a buffer to fit at least 'size' sizes and initializes them to
    // zero.
    inline BufferPtr allocateSizes(vector_size_t size, memory::MemoryPool *pool) {
        return AlignedBuffer::allocate<vector_size_t>(size, pool, 0);
    }
} // namespace kumo::pollux
